/*
 * AIEngine a new generation network intrusion detection system.
 *
 * Copyright (C) 2013-2023  Luis Campo Giralte
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Library General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Library General Public License for more details.
 *
 * You should have received a copy of the GNU Library General Public
 * License along with this library; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301, USA.
 *
 * Written by Luis Campo Giralte <luis.camp0.2009@gmail.com> 
 *
 */
#include "ModbusProtocol.h"
#include <iomanip>

namespace aiengine {

ModbusProtocol::ModbusProtocol():
	Protocol("Modbus", IPPROTO_TCP) {}

bool ModbusProtocol::check(const Packet &packet) {

	int length = packet.getLength();

	if (length >= header_size) {
		if ((packet.getSourcePort() == 502)||(packet.getDestinationPort() == 502)) {
			// setHeader(packet.getPayload());
			++total_valid_packets_;
			return true;
		}
	}
	++total_invalid_packets_;
	return false;
}

void ModbusProtocol::processFlow(Flow *flow) {

	CPUCycle cycles(&total_cpu_cycles_);
	int length = flow->packet->getLength();
	total_bytes_ += length;

	++total_packets_;

	if (length > header_size) {
		setHeader(flow->packet->getPayload());	
		if (ntohs(header_->length) >= sizeof(modbus_header)) {
			const modbus_header *hdr = reinterpret_cast<const modbus_header*>(header_->data);

			if (hdr->code == MB_CODE_READ_COILS ) {
				++total_read_coils_;
			} else if (hdr->code == MB_CODE_READ_DISCRETE_INPUTS ) {
				++total_read_discrete_inputs_;
			} else if (hdr->code == MB_CODE_READ_HOLDING_REGISTERS ) {
				++total_read_holding_registers_;
			} else if (hdr->code == MB_CODE_READ_INPUT_REGISTERS ) {
				++total_read_input_registers_;
			} else if (hdr->code == MB_CODE_WRITE_SINGLE_COIL ) {
				++total_write_single_coil_;
			} else if (hdr->code == MB_CODE_WRITE_SINGLE_REGISTER ) {
				++total_write_single_register_;
			} else if (hdr->code == MB_CODE_WRITE_MULTIPLE_COILS ) {
				++total_write_multiple_coils_;
			} else if (hdr->code == MB_CODE_WRITE_MULTIPLE_REGISTERS ) {
				++total_write_multiple_registers_;
			} else {
				++total_others_;
			}
		}
	}
}

void ModbusProtocol::statistics(std::basic_ostream<char> &out, int level, int32_t limit) const { 

	showStatisticsHeader(out, level);

	if (level > 3) {
		out << "\t" << "Total read coils:       " << std::setw(10) << total_read_coils_ << "\n"
			<< "\t" << "Total read dis inputs:  " << std::setw(10) << total_read_discrete_inputs_ << "\n"
			<< "\t" << "Total read hold regs:   " << std::setw(10) << total_read_holding_registers_ << "\n"
			<< "\t" << "Total read input regs:  " << std::setw(10) << total_read_input_registers_ << "\n"
			<< "\t" << "Total write single coil:" << std::setw(10) << total_write_single_coil_  << "\n"
			<< "\t" << "Total write multi coils:" << std::setw(10) << total_write_multiple_coils_ << "\n"
			<< "\t" << "Total write multi regs: " << std::setw(10) << total_write_multiple_registers_ << "\n"
			<< "\t" << "Total others:           " << std::setw(10) << total_others_ << std::endl;
	}
	if ((level > 5)and(flow_forwarder_.lock()))
		flow_forwarder_.lock()->statistics(out);
}

void ModbusProtocol::statistics(Json &out, int level) const {          

	showStatisticsHeader(out, level);

	if (level > 3) {
		out["read coils"] = total_read_coils_;
		out["read dis inputs"] = total_read_discrete_inputs_;
		out["read hold regs"] = total_read_holding_registers_;
		out["read input regs"] = total_read_input_registers_;
		out["write single coil"] = total_write_single_coil_;
		out["write multi coils"] = total_write_multiple_coils_;
		out["write multi regs"] = total_write_multiple_registers_; 
		out["others"] = total_others_;
	}
}

CounterMap ModbusProtocol::getCounters() const {
  	CounterMap cm;
 
        cm.addKeyValue("packets",total_packets_);
        cm.addKeyValue("bytes", total_bytes_);

	cm.addKeyValue("read coils", total_read_coils_ );
	cm.addKeyValue("read dis inputs", total_read_discrete_inputs_ );
	cm.addKeyValue("read hold regs", total_read_holding_registers_ );
	cm.addKeyValue("read input regs", total_read_input_registers_ );
	cm.addKeyValue("write single coil", total_write_single_coil_ );
	cm.addKeyValue("write multi coils", total_write_multiple_coils_ );
	cm.addKeyValue("write multi regs", total_write_multiple_registers_ );
	cm.addKeyValue("others", total_others_ );

        return cm;
}

void ModbusProtocol::resetCounters() {

	reset();

        total_read_coils_ = 0;
        total_read_discrete_inputs_ = 0;
        total_read_holding_registers_ = 0;
        total_read_input_registers_ = 0;
        total_write_single_coil_ = 0;
        total_write_single_register_ = 0;
        total_write_multiple_coils_ = 0;
        total_write_multiple_registers_ = 0;
        total_others_ = 0;
}

} // namespace aiengine
